#include "LowpolyProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"

void LowpolyProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();
	ubo.model = glm::mat4(1.f);
	ubo.scale = glm::vec3(1.f);
	ubo.scale /= mLocalScene->mVoxelResolution;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo ubodata = ubo;
			mUboBuffer = CustomBuffer::Builder()
				.SetData(&ubodata)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			GeometryGenerator geometryGenerator;
			GeometryGenerator::MeshData cubeData = geometryGenerator.CreateCube(1.f, 1.f, 1.f, 1);

			std::vector<LocalVertex> localVertexs;
			localVertexs.resize(cubeData.Vertices.size());
			for (uint32 i = 0; i < cubeData.Vertices.size(); ++i)
			{
				localVertexs[i].position = cubeData.Vertices[i].Position;
				localVertexs[i].normal = cubeData.Vertices[i].Normal;
			}

			mVertexBuffer = VertexBuffer::Builder()
				.SetData(localVertexs.data())
				.SetSize(localVertexs.size() * sizeof(LocalVertex))
				.Build();

			mIndexBuffer = IndexBuffer::Builder()
				.SetData(cubeData.Indices32.data())
				.SetSize(cubeData.Indices32.size() * sizeof(uint32))
				.SetIndexCount(cubeData.Indices32.size())
				.Build();

			RefCountPtr<RHITexture> voxelOccupancy = GetTexture("VoxelProcess", "voxelOccupancy");

			glm::vec3 voxelResolution = mLocalScene->mVoxelResolution;

			std::vector<uint32> voxelOccupancyData;
			voxelOccupancyData.resize(voxelResolution[0] * voxelResolution[1] * voxelResolution[2], 0u);
			voxelOccupancy->ReadMemory(voxelOccupancyData.data());

			std::vector<glm::vec3> voxelOffsets;
			for (uint32 i = 0; i < voxelResolution[2]; ++i)
			{
				for (uint32 j = 0; j < voxelResolution[1]; ++j)
				{
					for (uint32 k = 0; k < voxelResolution[0]; ++k)
					{
						int32 index = i * voxelResolution[1] * voxelResolution[0] + j * voxelResolution[0] + k;

						if (voxelOccupancyData[index])
						{
							voxelOffsets.push_back(glm::vec3(
								(float)k / voxelResolution[0],
								(float)j / voxelResolution[1],
								(float)i / voxelResolution[2]
							));
						}
					}
				}
			}

			assert(voxelOffsets.size() > 0);

			mInstanceSize = voxelOffsets.size();

			mInstanceBuffer = VertexBuffer::Builder()
				.SetData(voxelOffsets.data())
				.SetSize(voxelOffsets.size() * sizeof(glm::vec3))
				.Build();



			mDescriptor = RHIDescriptor::Builder()
				.PushShader<LowpolyVertexShader>()
				.PushShader<LowpolyFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();


			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build();

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void LowpolyProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();
	ubo.model = glm::mat4(1.f);
	ubo.scale = glm::vec3(1.f);
	ubo.scale /= mLocalScene->mVoxelResolution;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkDeviceSize offsets = 0;

			VkBuffer vertexBuffer = mVertexBuffer->GetBuffer();
			vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

			VkBuffer instaceBuffer = mInstanceBuffer->GetBuffer();
			vkCmdBindVertexBuffers(commandBuffer, 1, 1, &instaceBuffer, &offsets);

			VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

			vkCmdDrawIndexed(commandBuffer, mIndexBuffer->GetIndexCount(), mInstanceSize, 0, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void LowpolyProcess::OnResize()
{
	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build();

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
